import os
import re
import cv2
from typing import List, Tuple
from pathlib import Path
import chardet

class SubtitleProcessor:
    def __init__(self, subtitle_file: str, frames: List[str], merge_threshold: int = 1024):
        self.subtitle_file = subtitle_file
        self.frames = sorted(frames, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
        self.merge_threshold = merge_threshold
        self.subtitles = self._load_subtitles()
        self.time_segments = self._calculate_time_segments()

    def _load_subtitles(self) -> List[dict]:
        # 使用 pathlib 处理跨平台路径
        subtitle_file_path = Path(self.subtitle_file)

        # 检测文件编码
        with subtitle_file_path.open('rb') as f:
            raw_data = f.read()
            encoding = chardet.detect(raw_data)['encoding'] or 'utf-8'  # 默认使用 utf-8

        # 使用检测到的编码读取文件内容
        with subtitle_file_path.open('r', encoding=encoding) as f:
            content = f.read()

        # 统一换行符为 \n
        content = content.replace('\r\n', '\n').replace('\r', '\n')

        return self._split_text_into_chunks(content)

    def _split_text_into_chunks(self, input_data: str) -> List[dict]:
        output = []
        for line in input_data.strip().split('\n'):
            match = re.match(r'(\d{2}:\d{2}:\d{2})--(\d{2}:\d{2}:\d{2})\s+(.*)', line)
            if match:
                start_time, end_time, text = match.groups()
                output.append({
                    "timestamp": (self._time_to_seconds(start_time), self._time_to_seconds(end_time)),
                    "text": text
                })
        return output

    @staticmethod
    def _time_to_seconds(time_str: str) -> int:
        hours, minutes, seconds = map(int, time_str.split(':'))
        return hours * 3600 + minutes * 60 + seconds

    def _calculate_time_segments(self) -> List[Tuple[float, float]]:
        video = cv2.VideoCapture(self.frames[0])  # 使用第一帧来获取视频信息
        fps = video.get(cv2.CAP_PROP_FPS)
        video.release()

        time_segments = []
        for i, frame in enumerate(self.frames):
            frame_num = int(os.path.splitext(os.path.basename(frame))[0])
            start_time = frame_num / fps
            if i + 1 < len(self.frames):
                next_frame_num = int(os.path.splitext(os.path.basename(self.frames[i+1]))[0])
                end_time = (next_frame_num - 1) / fps
            else:
                end_time = start_time + 1  # 假设最后一帧持续1秒
            time_segments.append((start_time, end_time))
        return time_segments

    def process_subtitles(self) -> List[Tuple[str, str]]:
        processed_subtitles = []
        for i, (start_time, end_time) in enumerate(self.time_segments):
            filtered_subtitles = self._filter_subtitles(start_time, end_time)
            subtitle_text = self._format_subtitles(filtered_subtitles)
            if len(subtitle_text) < self.merge_threshold:
                # 合并小字幕
                if processed_subtitles:
                    processed_subtitles[-1] = (processed_subtitles[-1][0], 
                                               processed_subtitles[-1][1] + "\n" + subtitle_text)
                else:
                    processed_subtitles.append((self.frames[i], subtitle_text))
            else:
                processed_subtitles.append((self.frames[i], subtitle_text))
        return processed_subtitles

    def _filter_subtitles(self, start_time: float, end_time: float) -> List[dict]:
        return [chunk for chunk in self.subtitles 
                if chunk['timestamp'][0] < end_time and chunk['timestamp'][1] > start_time]

    @staticmethod
    def _format_subtitles(chunks: List[dict]) -> str:
        return "\n".join(chunk['text'] for chunk in chunks)

    @staticmethod
    def format_time(seconds: float) -> str:
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        seconds = int(seconds % 60)
        return f"{hours:02}:{minutes:02}:{seconds:02}"


# 使用示例
# subtitle_file = "path/to/subtitle.txt"
# frames = ["path/to/frame1.jpg", "path/to/frame2.jpg", ...]
# processor = SubtitleProcessor(subtitle_file, frames, merge_threshold=1024)
# processed_subtitles = processor.process_subtitles()

# for frame, subtitle in processed_subtitles:
#     print(f"Frame: {frame}")
#     print(f"Subtitle: {subtitle}\n")
